#!/usr/bin/env python
# coding=utf-8
print("正在初始化摄像头...")

import cv2
import os
import datetime

cap = cv2.VideoCapture(1)

print("初始化成功！")
 
# name = 'pear' #在这里改类的名字
 name = 'apple'
# name = 'banana'
# name = 'veg_leaf'
# name = 'cucamber' 
# name = 'ciga'
# name = 'bat'
# name = 'bottle'
# name = 'can'
# name = 'brick'

savedpath = './pictures/'+name
isExists = os.path.exists(savedpath)
if not isExists:
    os.makedirs(savedpath)
    print('path of %s is build' % (savedpath))
else:
    print('path of %s already exist and rebuild' % (savedpath))
print("按N键拍摄图片")
i=0
while(True):
    ret, frame = cap.read()
    gray = cv2.cvtColor(frame, 1)
    r,g,b = cv2.split(frame)
    
    r_avg = cv2.mean(r)[0]
    g_avg = cv2.mean(g)[0]
    b_avg = cv2.mean(b)[0]
    
    k = (r_avg + g_avg + b_avg) / 3
    kr = k / r_avg
    kg = k / g_avg
    kb = k / b_avg
    
    r = cv2.addWeighted(src1 = r, alpha = kr, src2 = 0, beta = 0, gamma = 0)
    g = cv2.addWeighted(src1 = g, alpha = kg, src2 = 0, beta = 0, gamma = 0)
    b = cv2.addWeighted(src1 = b, alpha = kb, src2 = 0, beta = 0, gamma = 0)
    
    balanced_frame = cv2.merge([r, g, b])

    rangeParam = 128
    roi_x = 350
    roi_w = 0
    roi_y = 250
    roi_h = 0
    roi = balanced_frame[roi_y - rangeParam:roi_y + roi_h + rangeParam, roi_x - rangeParam:roi_x + roi_w + rangeParam]
    balanced_frame = roi

    cv2.imshow('test',balanced_frame)
    now = datetime.datetime.now()
    now = now.strftime('%m-%d-%H-%M-%S')
    savedname = '/'+name+ '_' + str(i) + '_{0}' '.jpg'.format(now)
    if cv2.waitKey(1) & 0xFF == ord('n'):    #按N拍摄
        i += 1
        cv2.imwrite(savedpath + savedname, balanced_frame)
        cv2.namedWindow("Image")
        cv2.imshow("Image", balanced_frame)
        cv2.waitKey(0)
        cv2.destroyAllWindows()
 
cap.release()
cv2.destroyAllWindows()
